import torch
import torch.utils.data as data

import os
import sys
import importlib
import glob
from tqdm import tqdm
import time
from config import cfg
import torch.nn as nn

from trainframework import TrainFrameWork
from loss import dice_bce_loss
#from dataset import Dataset
#from dataset import Dataset_20

from showPreImg import show_pre_img
import pandas as pd
import numpy as np
from mean_error import MeanError
def test_net(epoch, solver, test_data_loader, save_path=None, test_mode=False):
    """

    :param epoch: 周期
    :param solver: 神经网络
    :param test_data_loader: 测试数据集
    :param test_mode: 测试模式 train 为False valid 为True
    :param save_path:  输出图片文件保存路径
    :return:
    """
    test_tic = time.time()
    test_data_loader_iter = iter(test_data_loader)
    accuracy= 0
    mean = 0
    pbar = tqdm(total=len(test_data_loader_iter))  # 进度条

    solver.eval()

    for step, (_, img, mask) in enumerate(test_data_loader_iter):
        solver.set_input(img, mask)

        acc, mask_pre = solver.test_batch()
        if test_mode:  # 如果为验证模式 打印图片
            show_pre_img(img, mask, mask_pre, step, save_path)
        mean_err_func = MeanError(mask, mask_pre)
        mean += mean_err_func.mean_ab_error()
        accuracy += acc

        pbar.update(1)
        pbar.set_description("Epoch: %d, Accuracy: %f ,Mean Error: %f" % (epoch, accuracy / (step + 1), mean / (step + 1.)))

    solver.train()
    # 关闭进度条
    pbar.close()
    time.sleep(0.1)  # 略微等待, 避免输出冲突
    sys.stdout.flush()
    # 计算本轮平均IoU
    accuracy /= len(test_data_loader_iter)
    # 输出测试结果
    print('Epoch: %3d        Accuracy: %.8f        elpase(s): %.3f' %
          (epoch, accuracy, time.time() - test_tic))

    return accuracy


if __name__ == "__main__":
    #viz = Visdom(env='lossa')
    # 导入配置 训练集路径
    TRAIN_PATH = cfg.TRAIN_PATH


    image_list = filter(lambda x: x.find('.bmdl') != -1, os.listdir(TRAIN_PATH))
    train_set = set(map(lambda x: x[:-5], image_list))

    image_list2 = filter(lambda x: x.find('.dat') != -1, os.listdir(TRAIN_PATH))
    train_set2 = set(map(lambda x: x[:-4], image_list2))

    train_set = train_set & train_set2
    train_list = list(train_set)
    # 导入配置 测试集路径
    TEST_PATH = cfg.TEST_PATH
    image_list = filter(lambda x: x.find('.bmdl') != -1, os.listdir(TEST_PATH))
    test_set = set(map(lambda x: x[:-5], image_list))

    image_list2 = filter(lambda x: x.find('.dat') != -1, os.listdir(TEST_PATH))
    test_set2 = set(map(lambda x: x[:-4], image_list2))

    test_set = test_set & test_set2
    test_list = list(test_set)

    # 初始化保存路径 若路径不存在则创建
    SAVE_PATH = cfg.SAVE_PATH
    if not (os.path.exists(SAVE_PATH)):
        os.makedirs(SAVE_PATH)
    # 更具config.ini 导入对应网络模型
    net = getattr(importlib.import_module(cfg.MODULE), cfg.MODEL)
    # 损失函数
    loss = nn.SmoothL1Loss(reduction='mean')  # nn.MSELoss()
    # 传入train Framework框架
    solver = TrainFrameWork(net, cfg.OPTIMIZER, loss, cfg.LR)



    ## 载入已有模型
    epoch = 1
    modelfiles = glob.glob('%s/model0*.model' % SAVE_PATH)
    modelfiles.sort()
    INIT_MODEL = cfg.INIT_MODEL
    if INIT_MODEL != "":
        solver.load(INIT_MODEL)
        print("Model %s loaded!" % INIT_MODEL);
        epoch = int(os.path.splitext(os.path.basename(INIT_MODEL))[0][5:]) + 1
    elif len(modelfiles) >= 1:
        solver.load(modelfiles[-1]);
        print("Model %s loaded from previous state!" % modelfiles[-1]);
        epoch = int(os.path.splitext(os.path.basename(modelfiles[-1]))[0][5:]) + 1


    ## 数据集
    batch_size = torch.cuda.device_count() * cfg.BATCH_SIZE_PER_CARD
    # color 参数的作用暂时不明白
    dataset_getattr = getattr(importlib.import_module(cfg.DATA_LOAD_MODULE), cfg.DATA_LOAD_MODEL)

    dataset = dataset_getattr(list(train_list), TRAIN_PATH, cfg.MASK_COLOR)

    # dataset = Dataset(list(train_list), TRAIN_PATH, cfg.MASK_COLOR)

    data_loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=cfg.DATA_LOADER_NUM_WORKERS
    )


    test_dataset = dataset_getattr(list(test_list), TEST_PATH, cfg.MASK_COLOR, eval=True)



    test_data_loader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=cfg.DATA_LOADER_NUM_WORKERS
    )

    if cfg.EVAL:
        test_net(epoch-1, solver, test_data_loader,  cfg.SAVE_PATH, cfg.EVAL)

        exit(0)

    ##　日志文件
    LOG_FILE_NAME = os.path.join(SAVE_PATH, 'log.txt')
    log_file = open(LOG_FILE_NAME, 'a+')
    LOSS_LOG_FILE_NAME = os.path.join(SAVE_PATH, 'loos_log.txt')
    loss_log_file = open(LOG_FILE_NAME, 'a+')

    no_optim = 0
    # 总训练周期数
    total_epoch = cfg.MAX_EPOCH
    # 最佳
    train_epoch_best_loss = 1000000000.

    best_iou = 0.0
    MAX_EPOCH = cfg.MAX_EPOCH
    # 测试间隔
    TEST_INTERVAL = cfg.TEST_INTERVAL

    while epoch <= MAX_EPOCH:

        epoch_tic = time.time()  # 周期计时器

        data_loader_iter = iter(data_loader) # 迭代对象
        train_epoch_loss = 0        # 周期平均损失
        pbar = tqdm(total=len(data_loader_iter))  # 进度条

        for ibatch, (_, x, y) in enumerate(data_loader_iter):
            # 传入x,y
            solver.set_input(x, y)
            train_loss = solver.optimize()
            train_epoch_loss += train_loss

            pbar.update(1)
            pbar.set_description("Epoch: %d, loss: %f" % (epoch, train_epoch_loss / (ibatch + 1)))

        # 关闭进度条
        pbar.close()
        time.sleep(0.1)  # 略微等待, 避免输出冲突
        sys.stdout.flush()

        # 计算本轮平均损失
        train_epoch_loss /= len(data_loader_iter)

        # 记录loss 文件便于之后用于loss分析
        loss_log_file.write('{},{},{}\n'.format(epoch, solver.lr, train_epoch_loss))

        # 测试
        if epoch % TEST_INTERVAL == 0:
            # 传入保存文件名
            solver.save(SAVE_PATH + "/model%09d.model" % epoch)

            accuracy = test_net(epoch, solver, test_data_loader)
            sys.stdout.flush()
            log_file.write("IT %d, LR %f, LOSS %f, Accuracy %f\n" % (epoch, solver.lr, train_epoch_loss, accuracy))
        else:
            log_file.write("IT %d, LR %f, LOSS %f\n" % (epoch, solver.lr, train_epoch_loss))


        # 判断是否是最佳loss 如果不是 计数加1 计数为10时结束训练
        if train_epoch_loss >= train_epoch_best_loss:
            no_optim += 1
        else:
            no_optim = 0
            train_epoch_best_loss = train_epoch_loss
            # solver.save('weights/'+NAME+'.th')
        if no_optim > 10:
            log_file.write('early stop at %d epoch' % epoch)
            print('early stop at %d epoch' % epoch)
            break
        if no_optim > 3:
            if solver.lr < 5e-7:
                break
            #solver.load('weights/'+NAME+'.th')

            solver.update_lr(0.2, factor=True)

        log_file.flush()
        epoch += 1

    log_file.close()